package edu.zju.gis.dbfg.model.etl;

import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.base.GeoItem;
import edu.zju.gis.dbfg.model.scala.MRasterizer2;
import geotrellis.raster.MultibandTile;
import geotrellis.raster.RasterExtent;
import geotrellis.vector.Extent;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.broadcast.Broadcast;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.io.WKTReader;
import org.opengis.referencing.FactoryException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.util.*;


/**
 * @author Hu
 * @date 2019/10/30
 * 矢量栅格化，输出为多波段tiff，demo2
 **/
public class demo2 {

  private static final Logger logger = LoggerFactory.getLogger(demo2.class);

  public static void main(String[] args) throws FactoryException {

    logger.info("Geotrellis AreaStatistic");
    // Setup environment
    logger.info("Setup Spark Context");
    SparkConf conf = new SparkConf();
    conf.setMaster("local[1]");
    conf.setAppName("lalala");
    conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
    // conf.registerKryoClasses();
    JavaSparkContext jsc = new JavaSparkContext(conf);

    List<String> files = new ArrayList<>();
    files.add("D:\\Work\\Study\\GeRenXueXi\\Code\\4_矢量瓦片\\data_宁波北仑_LCRA_polygon\\LCRA.wkt");

    int z = 20;

    PyramidConfig pc =  new PyramidConfig.PyramidConfigBuilder()
      .setBaseMapEnv(-180, 180, -90, 90)
      .setCrs(CRS.decode("EPSG:4326"))
      .setzLevelRange(4, 22)
      .build();

    Broadcast<PyramidConfig> pcBroad = jsc.broadcast(pc);
    Broadcast<Integer> zBroad = jsc.broadcast(z);

    JavaRDD<String> filesRDD = jsc.parallelize(files, files.size());
    JavaRDD<String> result = filesRDD.map(new Function<String, String>() {
      @Override
      public String call(String s) throws Exception {
        String msg = "Generate GeoTiff file success";
        BufferedReader bf = new BufferedReader(new FileReader(new File(s)));
        String line = bf.readLine();
        List<GeoItem> gis = new ArrayList<>();
        double xmin = 180d;
        double xmax = -180d;
        double ymin = 90d;
        double ymax = -90d;
        while (line != null) {
          if (line.trim().length() == 0) continue;
          GeoItem gi = new GeoItem();
          Map<String, Object> props = new HashMap<>();
          String[] fields = line.split("\t");
          String wkt = fields[fields.length-1];
          WKTReader reader = new WKTReader();
          Geometry geometry = reader.read(wkt);
          gi.setGeometry(geometry);
          Envelope e = geometry.getEnvelopeInternal();
          xmin = e.getMinX() < xmin ? e.getMinX():xmin;
          xmax = e.getMaxX() > xmax ? e.getMaxX():xmax;
          ymin = e.getMinY() < ymin ? e.getMinY():ymin;
          ymax = e.getMaxY() > ymax ? e.getMaxY():ymax;
          props.put("CC", fields[2]);
          gi.setProp(props);
          gi.setFid(fields[1]);
          gis.add(gi);
          line = bf.readLine();
        }

        Extent extent = new Extent(xmin, ymin, xmax, ymax);
        Double xresolution = pcBroad.getValue().getBaseMapEnv().getWidth()/Math.pow(2, zBroad.getValue());
        Double yresolution = pcBroad.getValue().getBaseMapEnv().getHeight()/Math.pow(2, zBroad.getValue());
        RasterExtent re = MRasterizer2.getRasterExtent(extent, xresolution, yresolution);
        MultibandTile mbt = MRasterizer2.rasterizeWithValue(gis, re);
        String outputPath = "D:\\Work\\Study\\GeRenXueXi\\Code\\4_矢量瓦片\\data_宁波北仑_LCRA_polygon\\LCRA.tiff";
        MRasterizer2.write(mbt, extent, outputPath);
        return msg;
      }
    });

    logger.info(String.format("Success: %d", result.count()));
  }

}
